import glob
import os

import numpy as np
import SimpleITK as sitk

segrap_task001_one_hot_label_names = {
    'Brain': 1,
    "BrainStem": 2,
    "Chiasm": 3,
    "TemporalLobe_L": 4,
    "TemporalLobe_R": 5,
    "TemporalLobe_Hippocampus_OverLap_L": 6,
    "TemporalLobe_Hippocampus_OverLap_R": 7,
    "Hippocampus_L": 8,
    "Hippocampus_R": 9,
    'Eye_L': 10,
    'Eye_R': 11,
    "Lens_L": 12,
    "Lens_R": 13,
    "OpticNerve_L": 14,
    "OpticNerve_R": 15,
    "MiddleEar_L": 16,
    "MiddleEar_R": 17,
    "IAC_L": 18,
    "IAC_R": 19,
    "MiddleEar_TympanicCavity_OverLap_L": 20,
    "MiddleEar_TympanicCavity_OverLap_R": 21,
    "TympanicCavity_L": 22,
    "TympanicCavity_R": 23,
    "MiddleEar_VestibulSemi_OverLap_L": 24,
    "MiddleEar_VestibulSemi_OverLap_R": 25,
    "VestibulSemi_L": 26,
    "VestibulSemi_R": 27,
    "Cochlea_L": 28,
    "Cochlea_R": 29,
    "MiddleEar_ETbone_OverLap_L": 30,
    "MiddleEar_ETbone_OverLap_R": 31,
    "ETbone_L": 32,
    "ETbone_R": 33,
    "Pituitary": 34,
    "OralCavity": 35,
    "Mandible_L": 36,
    "Mandible_R": 37,
    "Submandibular_L": 38,
    "Submandibular_R": 39,
    "Parotid_L": 40,
    "Parotid_R": 41,
    "Mastoid_L": 42,
    "Mastoid_R": 43,
    "TMjoint_L": 44,
    "TMjoint_R": 45,
    "SpinalCord": 46,
    "Esophagus": 47,
    "Larynx": 48,
    "Larynx_Glottic": 49,
    "Larynx_Supraglot": 50,
    "Larynx_PharynxConst_OverLap": 51,
    "PharynxConst": 52,
    "Thyroid": 53,
    "Trachea": 54
}


segrap_task002_one_hot_label_names = {
    "GTVp": 1,
    "GTVnd": 2}


segrap_subset = {
    'Brain': [1, 2, 3, 4, 5, 6, 7, 8, 9],
    "BrainStem": 2,
    "Chiasm": 3,
    "TemporalLobe_L": [4, 6],
    "TemporalLobe_R": [5, 7],
    "Hippocampus_L": [8, 6],
    "Hippocampus_R": [9, 7],
    'Eye_L': [10, 12],
    'Eye_R': [11, 13],
    "Lens_L": 12,
    "Lens_R": 13,
    "OpticNerve_L": 14,
    "OpticNerve_R": 15,
    "MiddleEar_L": [18, 16, 20, 24, 28, 30],
    "MiddleEar_R": [19, 17, 21, 25, 29, 31],
    "IAC_L": 18,
    "IAC_R": 19,
    "TympanicCavity_L": [22, 20],
    "TympanicCavity_R": [23, 21],
    "VestibulSemi_L": [26, 24],
    "VestibulSemi_R": [27, 25],
    "Cochlea_L": 28,
    "Cochlea_R": 29,
    "ETbone_L": [32, 30],
    "ETbone_R": [33, 31],
    "Pituitary": 34,
    "OralCavity": 35,
    "Mandible_L": 36,
    "Mandible_R": 37,
    "Submandibular_L": 38,
    "Submandibular_R": 39,
    "Parotid_L": 40,
    "Parotid_R": 41,
    "Mastoid_L": 42,
    "Mastoid_R": 43,
    "TMjoint_L": 44,
    "TMjoint_R": 45,
    "SpinalCord": 46,
    "Esophagus": 47,
    "Larynx": [48, 49, 50, 51],
    "Larynx_Glottic": 49,
    "Larynx_Supraglot": 50,
    "PharynxConst": [51, 52],
    "Thyroid": 53,
    "Trachea": 54}


def nii2array(path):
    mask_itk_ref = sitk.ReadImage(path)
    mask_arr_ref = sitk.GetArrayFromImage(mask_itk_ref)
    return mask_arr_ref


def merge_multi_class_to_one(input_arr, classes_index=None):
    new_arr = np.zeros_like(input_arr)
    for cls_ind in classes_index:
        new_arr[input_arr == cls_ind] = 1
    return new_arr


def convert_one_hot_label_to_multi_organs(ont_hot_label_path, save_fold):
    for organ in segrap_subset.keys():
        ont_hot_label_arr = nii2array(ont_hot_label_path)
        ont_hot_label_itk = sitk.ReadImage(ont_hot_label_path)
        if type(segrap_subset[organ]) is list:
            new_arr = merge_multi_class_to_one(
                ont_hot_label_arr, segrap_subset[organ])
        else:
            new_arr = np.zeros_like(ont_hot_label_arr)
            new_arr[ont_hot_label_arr == segrap_subset[organ]] = 1
        new_itk = sitk.GetImageFromArray(new_arr)
        new_itk.CopyInformation(ont_hot_label_itk)
        sitk.WriteImage(new_itk, "{}/{}.nii.gz".format(save_fold, organ))


if __name__ == "__main__":
    for patient in glob.glob("../nnUNet_InfersTs/Task001/*"):
        new_path = "../Submission_Task001/{}".format(
            patient.split("/")[-1].replace(".nii.gz", ""))
        if os.path.exists(new_path):
            pass
        else:
            os.mkdir(new_path)
            convert_one_hot_label_to_multi_organs(patient, new_path)
